import pandas as pd
from prefect import get_run_logger

from application.db import get_database_connection
from application.db.mysql_db.nsfc.NsfcInfoAttachmentList import NsfcInfoAttachmentList
from application.db.mysql_db.nsfc.NsfcInfoList import NsfcInfoList
from application.db.mysql_db.nsfc.NsfcInfoSectionList import NsfcInfoSectionList

from application.tasks.base_task import BaseTask


class DataStorageTask(BaseTask):
    @get_database_connection("default").atomic()  # 保证事务
    def run(self, df: pd.DataFrame):

        df['info_name'] = df['information_name'].apply(lambda x: x['zh'])  # 添加中文标题字段

        information_list = df[
            ['information_id', 'info_type_id', 'apply_code', 'source_id', 'info_name', 'original_link',
             'publish_date']].dropna(axis=1, how='all').dropna(
            how='all')  # 添加信息表字段

        information_section = df['info_section'].explode().apply(pd.Series).dropna(axis=1, how='all').dropna(
            how='all')  # 添加章节表字段

        info_attachment = df['info_attachment'].explode().apply(pd.Series).dropna(axis=1, how='all').dropna(
            how='all')  # 添加附件表字段

        # 保证所有待入库的 DataFrame 均无 NaN

        # 如果 information_list 不为空才入库
        if not information_list.empty:
            NsfcInfoList.insert_many(self.clean_dataframe(information_list).to_dict(orient="records")).execute()

        if not information_section.empty:
            NsfcInfoSectionList.insert_many(self.clean_dataframe(information_section).to_dict(orient="records")).execute()

        if not info_attachment.empty:
            NsfcInfoAttachmentList.insert_many(self.clean_dataframe(info_attachment).to_dict(orient="records")).execute()

        return df
